package cn.swing.main.srv.cv.tracker;

import cn.swing.main.srv.cv.model.TennisPoint;
import cn.swing.main.srv.cv.model.Trajectory;
import cn.swing.main.srv.cv.utils.HungarianAlgorithm;
import cn.swing.main.srv.cv.utils.TrajectoryMerger;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.jfree.data.json.impl.JSONArray;
import org.opencv.core.Core;
import org.opencv.core.CvType;
import org.opencv.core.Mat;
import org.opencv.core.Point;
import org.opencv.core.Scalar;
import org.opencv.core.Size;
import org.opencv.highgui.HighGui;
import org.opencv.imgproc.Imgproc;
import org.opencv.videoio.VideoCapture;
import org.opencv.videoio.VideoWriter;
import org.opencv.videoio.Videoio;
import org.springframework.util.CollectionUtils;

@Slf4j
public class AdvancedTracker {

    /**
     * 每帧小于{x}距离的点才关联轨迹
     */
    public static final double MAX_DISTANCE_PER_FRAME = 30;

    /**
     * MAX_MISSED_FRAMES帧之内最大距离小于{n}的点才关联轨迹
     */
    public static final double MAX_DISTANCE = 100;

    /**
     * 合并两个轨迹时，最大距离不能超过{n}
     */
    public static final double MERGE_MAX_DISTANCE = 200;

    /**
     * 合并轨迹时，两段轨迹帧差不能超过{x}
     */
    public static final double MERGE_TRAJ_FRAME_GAP_THRESHOLD = 10;

    /**
     * 合并轨迹时，两段轨迹重叠不能超过{x}
     */
    public static final double MERGE_TRAJ_FRAME_OVERLAP_THRESHOLD = 3;

    /**
     * 可丢失的轨迹帧数
     */
    private static final int MAX_MISSED_FRAMES = 8;

    /**
     * 只保存连续长度大于{n}的轨迹
     */
    private static final int SAVE_FRAME_POINT_NUMBER = 5;

    /**
     * 当前正在活跃的轨迹
     */
    List<Trajectory> trajectories = new ArrayList<>();

    Map<Integer, List<TennisPoint>> allFramePoints = new HashMap<>();

    int startFrameNumber = 0;

    int endFrameNumber = 0;

    private int maxTrajectoryId = 0;

    /**
     * 最终临时保存的所有轨迹
     */
    List<Trajectory> savedTrajectories = new ArrayList<>();

    public long processFrame(List<TennisPoint> framePoints, int frameNum) {
        if (startFrameNumber == 0) {
            startFrameNumber = frameNum;
        }
        long startTime = System.currentTimeMillis();
        this.endFrameNumber = frameNum;

        if (CollectionUtils.isEmpty(framePoints) || framePoints.size() > 100) {
            System.out.println("当前帧点数:" + framePoints.size() + "，跳过处理");
            return System.currentTimeMillis() - startTime;
        }
        // 先保存到所有的点中
        allFramePoints.put(frameNum, framePoints);
        // 步骤1：清理丢失轨迹并保存有价值的
        removeDeadTrajectories(frameNum, trajectories, savedTrajectories,
                MAX_MISSED_FRAMES, SAVE_FRAME_POINT_NUMBER, false);

        processPoints(trajectories, framePoints, frameNum);

        return System.currentTimeMillis() - startTime;

    }

    private void processPoints(List<Trajectory> trajectories,
            List<TennisPoint> framePoints, int frameNum) {
        if (!CollectionUtils.isEmpty(trajectories)) {
            // 步骤3：构建代价矩阵
            double[][] costMatrix = buildCostMatrix(trajectories, framePoints);

            // 步骤4：匈牙利算法匹配
            int[][] assignments = new HungarianAlgorithm(costMatrix).matchTracks();

            if (assignments != null) {
                // 步骤5：更新匹配轨迹
                updateTrajectories(trajectories, framePoints, assignments);
                // 步骤6：处理未匹配点（创建新轨迹）
                createNewTrajectories(trajectories, framePoints, assignments);
            }
        } else {
            // 步骤0 直接创建新的轨迹
            for (TennisPoint p : framePoints) {
                trajectories.add(new Trajectory(p, ++maxTrajectoryId));
            }
        }
    }

    private double[][] buildCostMatrix(List<Trajectory> trajectories, List<TennisPoint> points) {
        // 获取轨迹列表的大小
        int m = trajectories.size();
        // 获取点列表的大小
        int n = points.size();
        // 创建一个m行n列的矩阵，用于存储每个轨迹与每个点之间的距离
        double[][] matrix = new double[m][n];

        // 遍历所有轨迹
        for (int i = 0; i < m; i++) {
            // 获取当前轨迹
            Trajectory traj = trajectories.get(i);
            // 获取当前轨迹的状态估计向量的前两个元素（通常是x和y坐标）
            if (traj.getPoints().size() <= 3) {
                // 计算最后一个点和所有点的欧氏距离
                final TennisPoint last = traj.last();
                for (int j = 0; j < n; j++) {
                    // 获取当前点
                    TennisPoint p = points.get(j);
                    // 1. 计算位置特征的欧氏距离
                    double dx = last.getX() - p.getX();
                    double dy = last.getY() - p.getY();
                    double xyCost = Math.sqrt(dx * dx + dy * dy);
                    matrix[i][j] = xyCost;
                }
            } else {
                // 预测轨迹的下一个位置
                //                traj.getFilter().getFilter().predict();
                //                RealVector last = traj.getFilter().getFilter().getStateEstimationVector()
                //                .getSubVector(0, 4);
                final TennisPoint last = traj.last();
                final TennisPoint beforeLast = traj.getPoints().get(traj.getPoints().size() - 2);
                double preX = last.getVx() * 2 - beforeLast.getVx() + last.getX();
                double preY = last.getVy() * 2 - beforeLast.getVy() + last.getY();

                // 计算预测位置与所有点的欧氏距离
                for (int j = 0; j < n; j++) {
                    // 获取当前点
                    TennisPoint p = points.get(j);
                    // 1. 计算位置特征的欧氏距离
                    double dx = preX - p.getX();
                    double dy = preY - p.getY();
                    double xyCost = Math.sqrt(dx * dx + dy * dy);
                    matrix[i][j] = xyCost;
                }
            }

        }
        // 返回计算好的距离矩阵
        return matrix;
    }

    private void updateTrajectories(List<Trajectory> trajectories, List<TennisPoint> points, int[][] assignments) {
        for (int[] pair : assignments) {
            if (pair == null) {
                continue;
            }
            int trajIdx = pair[0];
            int pointIdx = pair[1];

            if (trajIdx >= 0 && pointIdx >= 0
                    // 因为代价矩阵会补零，因此这儿要注意可能出匹配出来超出的点的下标
                    && pointIdx < points.size()
                    // 因为代价矩阵会补零，因此这儿要注意可能出匹配出来超出的轨迹下标
                    && trajIdx < trajectories.size()) {
                Trajectory traj = trajectories.get(trajIdx);
                TennisPoint prev = traj.getPoints().get(traj.getPoints().size() - 1);

                TennisPoint curr = points.get(pointIdx);
                int frameDiff = curr.getFrameNumber() - prev.getFrameNumber();
                if (frameDiff == 0) {
                    return;
                }
                // 通过欧式距离判定
                double distance = Math.sqrt(Math.pow(curr.getX() - prev.getX(), 2)
                        + Math.pow(curr.getY() - prev.getY(), 2));
                // 首先距离超过了MAX_DISTANCE_PER_FRAME 再按后面的判断
                if (distance > MAX_DISTANCE_PER_FRAME
                        && distance > Math.min(prev.getVx() == 0 && prev.getVy() == 0 ?
                                MAX_DISTANCE_PER_FRAME * frameDiff :
                                // 放大3倍 避免误判
                                Math.sqrt(prev.getVx() * prev.getVx() + prev.getVy() * prev.getVy()) * 3 * frameDiff,
                        MAX_DISTANCE)) {
                    // 将当前轨迹的匹配点设为-1，表示该轨迹没有处理，后续会新建一条轨迹
                    pair[1] = -1;
                } else {
                    traj.setMissedFrames(0);
                    traj.setEndFrameNumber(curr.getFrameNumber());
                    traj.add(curr);
                    //                    try {
                    //                        traj.getFilter().update(curr);
                    //                    } catch (Exception e) {
                    //                        log.error("轨迹更新异常, {}", curr, e);
                    //                    }
                }
            }
        }
    }

    private void createNewTrajectories(List<Trajectory> trajectories, List<TennisPoint> points, int[][] assignments) {
        Set<Integer> matchedPoints = new HashSet<>();
        for (int[] pair : assignments) {
            if (pair == null) {
                continue;
            }
            if (pair[1] >= 0) {
                matchedPoints.add(pair[1]);
            }
        }

        for (int j = 0; j < points.size(); j++) {
            if (!matchedPoints.contains(j)) {
                trajectories.add(new Trajectory(points.get(j), ++maxTrajectoryId));
            }
        }
    }

    private void removeDeadTrajectories(int frameNum, List<Trajectory> trajectories,
            List<Trajectory> savedTrajectories, int missedFrames, int savePointNumber, boolean flush) {
        List<Trajectory> removedSave = new ArrayList<>();
        Iterator<Trajectory> it = trajectories.iterator();
        while (it.hasNext()) {
            Trajectory traj = it.next();
            // 计算轨迹丢失帧数
            traj.setMissedFrames(frameNum > traj.getEndFrameNumber() ? frameNum - traj.getEndFrameNumber() - 1 : 0);
            if (traj.getMissedFrames() > missedFrames) {
                // 将清理的轨迹进行甄别保留：连续长度大于10的轨迹
                if (traj.getPoints().size() > savePointNumber) {
                    // 保留轨迹
                    removedSave.add(traj);
                }
                it.remove();
            } else if (flush) {
                // 全部清理
                if (traj.getPoints().size() > savePointNumber) {
                    // 保留轨迹
                    removedSave.add(traj);
                }
                it.remove();
            }
        }
        if (!CollectionUtils.isEmpty(removedSave)) {
            savedTrajectories.addAll(removedSave);
        }
    }

    public void debugShow(boolean debug, boolean outputFlag, String path, double startFrame, double stopFrame) {
        if (!outputFlag) {
            return;
        }
        VideoCapture capture = new VideoCapture(path);
        if (!capture.isOpened()) {
            System.out.println("Error: Unable to open video file.");
            return;
        }
        // 初始化 VideoWriter outputPath是path文件名加一个out结尾
        // 获取capture.getFileName
        String pathName = path.substring(path.lastIndexOf("\\") + 1);
        // 在pathName前面加一个out
        String outputPath = "D:/swing/output/" + pathName;
        double fps = capture.get(Videoio.CAP_PROP_FPS);
        // 获取视频的宽度和高度
        int frameWidth = (int) capture.get(Videoio.CAP_PROP_FRAME_WIDTH);
        int frameHeight = (int) capture.get(Videoio.CAP_PROP_FRAME_HEIGHT);
        VideoWriter videoWriter = new VideoWriter(
                outputPath, // 输出文件路径
                VideoWriter.fourcc('m', 'p', '4', 'v'), // 编码器
                fps, // 帧率
                new Size(frameWidth, frameHeight) // 帧大小
        );

        if (!videoWriter.isOpened()) {
            System.out.println("Error: Unable to open video writer." + outputPath);
            return;
        }
        List<TennisPoint> tracePoints = new ArrayList<>();
        for (Trajectory traj : savedTrajectories) {
            tracePoints.addAll(traj.getPoints());
        }
        // 将轨迹中所有的点，按帧号分组
        final Map<Integer, List<TennisPoint>> tracePointGroup = tracePoints.stream()
                .collect(Collectors.groupingBy(TennisPoint::getFrameNumber));

        Mat frame = new Mat();
        int frameCount = 0;
        while (capture.read(frame)) {
            frameCount++;
            if (frameCount > stopFrame) {
                break;
            }
            if (frameCount < startFrame) {
                continue;
            }
            // 绘制轨迹点 红色
            final List<TennisPoint> tennisPoints = tracePointGroup.get(frameCount);
            if (!CollectionUtils.isEmpty(tennisPoints)) {
                for (TennisPoint point : tennisPoints) {
                    final Scalar scalar = new Scalar(0, 0, 125 + Math.abs(point.getTrajectoryId().hashCode() % 130));
                    Imgproc.rectangle(frame,
                            new Point(point.getX() - point.getWidth() / 2,
                                    point.getY() - point.getHeight() / 2),
                            new Point(point.getX() + point.getWidth() / 2,
                                    point.getY() + point.getHeight() / 2),
                            scalar, 2);
                    // 在框框的右边绘制轨迹id
                    Imgproc.putText(frame, point.getTrajectoryId() + " - " + point.getFrameNumber(),
                            new Point(point.getX() + point.getWidth(),
                                    point.getY()),
                            Imgproc.FONT_HERSHEY_SIMPLEX, 0.8,
                            scalar, 2);
                }
            }
            if (debug) {
                // 绘识别点 绿色
                final List<TennisPoint> allPoints = allFramePoints.get(frameCount);
                if (!CollectionUtils.isEmpty(allPoints)) {
                    for (TennisPoint point : allPoints) {
                        Imgproc.rectangle(frame,
                                new Point(point.getX() - point.getWidth() / 2,
                                        point.getY() - point.getHeight() / 2),
                                new Point(point.getX() + point.getWidth() / 2,
                                        point.getY() + point.getHeight() / 2),
                                new Scalar(255, 0, 0), 4);
                    }
                }
                HighGui.imshow("Tennis Ball Tracking", frame);
                final int i = HighGui.waitKey(1);
            }
            // 写入处理后的帧到视频文件
            videoWriter.write(frame);
        }
        frame.release();
        // 释放资源
        videoWriter.release();
        capture.release();
    }

    public void flushAndMerge() {
        if (!CollectionUtils.isEmpty(trajectories)) {
            removeDeadTrajectories(endFrameNumber, trajectories, savedTrajectories,
                    MAX_MISSED_FRAMES, SAVE_FRAME_POINT_NUMBER, true);
        }
        savedTrajectories.sort(Comparator.comparingInt(Trajectory::getId));
        // 打印merge前
        printTrajectories(savedTrajectories);
        // 前置清理，非网球的轨迹
        removeBeforeMerge();
        if (!CollectionUtils.isEmpty(savedTrajectories)) {
            // 从前往后merge
            List<Trajectory> toRemove = new ArrayList<>();
            for (int i = savedTrajectories.size() - 2; i > 0; i--) {
                Trajectory curr = savedTrajectories.get(i);
                for (int j = savedTrajectories.size() - 1; j > 0; j--) {
                    Trajectory next = savedTrajectories.get(j);
                    if (validateConnection(curr, next)) {
                        System.out.println("轨迹合并：" + curr.getId() + " - " + next.getId());
                        curr.merge(next);
                        toRemove.add(next);
                        break;
                    }
                }
            }
            savedTrajectories.removeAll(toRemove);
            System.out.println("merge完成，剩余轨迹数：" + savedTrajectories.size());
            // 合并后 清理无用的轨迹 长度低于10的，y轴整体变化不大的。 使用迭代器
            removeAfterMerge();
        }
        System.out.println("清理完成，剩余轨迹数：" + savedTrajectories.size());
        // 打印merge后
        printTrajectories(savedTrajectories);
        // 再执行一次匈牙利匹配
        savedTrajectories = finalMergeTrajectories(savedTrajectories);
        // 打印merge后
        printTrajectories(savedTrajectories);
    }

    private void removeBeforeMerge() {
        final Iterator<Trajectory> iterator = savedTrajectories.iterator();
        while (iterator.hasNext()) {
            Trajectory next = iterator.next();
            if (!isTennisTrajectory(next)) {
                iterator.remove();
            }
        }
    }

    private void removeAfterMerge() {
        final Iterator<Trajectory> iterator = savedTrajectories.iterator();
        while (iterator.hasNext()) {
            Trajectory t = iterator.next();
            if (t.getPoints().size() < 10) {
                iterator.remove();
            } else {
                // 删除轨迹中，累计y轴变化小于100像素的轨迹
                if (t.getPoints().stream().mapToDouble(e -> Math.abs(e.getVy())).sum() < 50) {
                    iterator.remove();
                } else if (t.getPoints().stream().mapToDouble(TennisPoint::getX).max().getAsDouble() -
                        t.getPoints().stream().mapToDouble(TennisPoint::getX).min().getAsDouble() <
                        150 &&
                        t.getPoints().stream().mapToDouble(TennisPoint::getY).max().getAsDouble() -
                                t.getPoints().stream().mapToDouble(TennisPoint::getY).min()
                                        .getAsDouble() < 100) {
                    // 删除轨迹中，x y 在固定一个区域的。
                    iterator.remove();
                }
            }
        }
    }

    private void printTrajectories(List<Trajectory> trajectoryList) {
        // 打印当前跟踪状态
        System.out.println("活跃轨迹数: " + trajectoryList.size());
        trajectoryList.forEach(t ->
                System.out.printf("轨迹ID=%d, 轨迹长度=%d, 开始帧=%d, 结束帧=%d, 漏点数量=%d 点列表=(%s) %n",
                        t.getId(),
                        t.getEndFrameNumber() - t.getStartFrameNumber(),
                        t.getStartFrameNumber(),
                        t.getEndFrameNumber(),
                        t.getEndFrameNumber() - t.getStartFrameNumber() - t.getPoints().size(),
                        JSONArray.toJSONString(t.getPoints())));
    }

    private boolean validateConnection(Trajectory t1, Trajectory t2) {
        if (t1.getEndFrameNumber() >= t2.getStartFrameNumber() + MERGE_TRAJ_FRAME_OVERLAP_THRESHOLD) {
            // 帧号不能重叠超过3帧
            return false;
        }
        if (t2.getStartFrameNumber() - t1.getEndFrameNumber() > MERGE_TRAJ_FRAME_GAP_THRESHOLD) {
            // 帧号不能丢失超过{MERGE_TRAJ_FRAME_DIFF_THRESHOLD}帧
            return false;
        }

        // 规则2：运动位置距离不能太远
        TennisPoint lastpoint = t1.last();
        TennisPoint firstpoint = t2.first();
        // 计算帧差
        int frameDiff = firstpoint.getFrameNumber() - lastpoint.getFrameNumber();
        double error = distance(new double[] {lastpoint.getX(), lastpoint.getY()},
                new double[] {firstpoint.getX(), firstpoint.getY()});

        // 首先距离超过了MAX_DISTANCE_PER_FRAME 再按后面的判断
        // 下面的参数和tracker 中设置的参数部分不一致，更宽松，因为两段轨迹之间要求不是特别严格
        if (error > MAX_DISTANCE_PER_FRAME
                && error > Math.min(lastpoint.getVx() == 0 && lastpoint.getVy() == 0 ?
                        MAX_DISTANCE_PER_FRAME * frameDiff :
                        // 放大3倍，避免速度变化过大
                        Math.sqrt(lastpoint.getVx() * lastpoint.getVx() + lastpoint.getVy() * lastpoint.getVy()) * 3 * frameDiff,
                // 这儿不应该用MAX_DISTANCE，因为两个轨迹合并的时候，长度超过了MAX_MISSED_FRAMES，所以这儿应该用每帧的长度，乘以帧差，避免速度变化过大
                MERGE_MAX_DISTANCE)) {
            return false;
        }
        return true;
    }

    private double distance(double[] p1, double[] p2) {
        double dx = p1[0] - p2[0];
        double dy = p1[1] - p2[1];
        return Math.sqrt(dx * dx + dy * dy);
    }

    public boolean isTennisTrajectory(Trajectory trajectory) {
        // 运动特征分析
        final double[] pm = hasParabolicMotion(trajectory.getPoints());
        if (pm != null && pm[0] > 0.4) {
            return false;
        }
        System.out.println(trajectory.getId() + " 轨迹加速度：" + hasValidVerticalAcceleration(trajectory.getPoints()));
        return true;
    }

    private double[] hasParabolicMotion(List<TennisPoint> points) {
        if (points.size() <= 5) {
            return null;
        }
        // 按顺序统计vx vy的正负变化次数
        int[] vc = new int[points.size() - 1];
        // 排除轨迹中第一个点，vx vy永远为0
        for (int i = 1; i < points.size(); i++) {
            // 根据vx vy的正负，设置vc[i]的值。都为正，设为象限1，vx负vy正，设为象限2，都为负，设为象限3，vx正vy负，设为象限4
            if (points.get(i).getVx() >= 0 && points.get(i).getVy() >= 0) {
                vc[i - 1] = 1;
            } else if (points.get(i).getVx() < 0 && points.get(i).getVy() >= 0) {
                vc[i - 1] = 2;
            } else if (points.get(i).getVx() < 0 && points.get(i).getVy() < 0) {
                vc[i - 1] = 3;
            } else {
                vc[i - 1] = 4;
            }
        }
        double changeCount = 0;
        for (int i = 1; i < vc.length; i++) {
            if (vc[i] != vc[i - 1]) {
                changeCount++;
            }
        }
        // 如果变化次数小于0.4，则认为是运动轨迹
        return new double[] {changeCount / vc.length, changeCount, vc.length};
    }

    /**
     * 这段代码通过线性回归拟合轨迹点的y坐标值，计算出加速度，并判断该加速度是否符合网球轨迹的特征。具体步骤包括：
     * 检查点的数量是否足够。
     * 提取时间数组和y值数组。
     * 创建并填充矩阵X和Y。
     * 使用Core.solve求解线性方程组。
     * 提取二次项系数并计算加速度。
     * 将加速度从像素单位转换为m/s²。
     * 判断加速度是否在合理范围内。
     *
     * @param points
     * @return
     */
    private double hasValidVerticalAcceleration(List<TennisPoint> points) {
        if (points.size() < 3) {
            return 0;
        }

        double[] times = new double[points.size()];
        double[] yValues = new double[points.size()];

        for (int i = 0; i < points.size(); i++) {
            times[i] = i; // Assuming equal time intervals for simplicity
            yValues[i] = points.get(i).getY();
        }

        Mat X = new Mat(times.length, 3, CvType.CV_64F);
        Mat Y = new Mat(times.length, 1, CvType.CV_64F);

        for (int i = 0; i < times.length; i++) {
            X.put(i, 0, 1);
            X.put(i, 1, times[i]);
            X.put(i, 2, times[i] * times[i]);
            Y.put(i, 0, yValues[i]);
        }

        Mat coefficients = new Mat();
        Core.solve(X, Y, coefficients, Core.DECOMP_SVD);

        double a = coefficients.get(2, 0)[0]; // Coefficient of t^2
        double g = 2 * a; // Convert to acceleration

        // Convert pixel acceleration to m/s^2 (assuming 1 pixel = 1 meter for simplicity)
        return g * 9.8; // Assuming 1 pixel = 1 meter
    }

    /**
     * 最终将轨迹的点放到一个轨迹中，并根据kalman滤波+匈牙利匹配判断，是否为网球轨迹
     * 重叠帧数 ≤10 且 断档帧数 ≤20
     *
     * @param trajectories
     * @return
     */
    public List<Trajectory> finalMergeTrajectories(List<Trajectory> trajectories) {
        Map<Integer, List<TennisPoint>> pointsFrame = new HashMap<>();
        for (int i = 0; i < trajectories.size(); i++) {
            for (int j = 0; j < trajectories.get(i).getPoints().size(); j++) {
                TennisPoint p = trajectories.get(i).getPoints().get(j);
                pointsFrame.computeIfAbsent(p.getFrameNumber(), f -> new ArrayList<>()).add(p);
            }
        }
        List<Trajectory> finalTrajectories = new ArrayList<>();
        List<Trajectory> saveTrajectories = new ArrayList<>();
        for (int i = startFrameNumber; i <= endFrameNumber; i++) {
            if (pointsFrame.containsKey(i)) {
                // 步骤1：清理丢失轨迹并保存有价值的
                removeDeadTrajectories(i, finalTrajectories, saveTrajectories,
                        MAX_MISSED_FRAMES * 3, SAVE_FRAME_POINT_NUMBER, false);
                processPoints(finalTrajectories, pointsFrame.get(i), i);
            }
        }
        saveTrajectories.addAll(finalTrajectories);
        // 合并这些轨迹
        TrajectoryMerger merger = new TrajectoryMerger();
        return merger.mergeTrajectories(saveTrajectories);
    }

}
